package nlp;

import java.io.IOException;
import java.util.List;

import bin.Configuration;
import edu.stanford.nlp.ie.AbstractSequenceClassifier;
import edu.stanford.nlp.ie.crf.CRFClassifier;
import edu.stanford.nlp.ling.CoreAnnotations.AnswerAnnotation;
import edu.stanford.nlp.ling.CoreLabel;

public class NamedEntityRecognizer {

	private AbstractSequenceClassifier<CoreLabel> classifier;

	@SuppressWarnings("unchecked")
    public NamedEntityRecognizer() {
		classifier = CRFClassifier.getClassifierNoExceptions(Configuration.ner_model_location);
	}
	
	public String tagText(String text) throws IOException {	     
		List<List<CoreLabel>> out = classifier.classify(text);		
		StringBuilder taggedText = new StringBuilder();
        StringBuilder name = new StringBuilder();
        String type = null;
        int start = -1;
        int end = -1;
        int last = 0;
		for (List<CoreLabel> sentence : out) {
			
			for (CoreLabel word : sentence) {
			
				//System.out.print(word + " ");
				//System.out.println(word.get(AnswerAnnotation.class));
								
				if (!word.get(AnswerAnnotation.class).equals("O")) {					
					if (name.length() > 0) {
						name.append(" ");
					}
					
					if (word.word().equalsIgnoreCase("-LRB-")) name.append("(");
					else if (word.word().equalsIgnoreCase("-RRB-")) name.append(")");
					else name.append(word.word());					
					if (type == null) type = word.get(AnswerAnnotation.class);
					if (start < 0) start = word.beginPosition();
					end = word.endPosition();
					
				} else if (name.length() > 0) {
				    int auxOffset = text.substring(0,start+1).split("\n+").length - 1;				    
				    int realStart = start + 1 - auxOffset;
				    int realEnd = end + 1 - auxOffset;
				    			    
				    //System.out.println("Name: " + name.toString() + " | Type: " + type + " | Start: " + realStart + " | End: " + realEnd);
				    
				    // build the tags: <LOC>,</LOC>, <ORG>,</ORG> or <PER>,</PER>
				    String s_tag = "<"+type.split("-")[1]+">";
				    String e_tag = "</"+type.split("-")[1]+"> ";
				    
				    // append previous text + s_tag + entity + end_tag				    
				    String textToAppend = text.substring(last,start);
				    taggedText.append(textToAppend);				    
				    taggedText.append(s_tag + name.toString() + e_tag);
				    last = realEnd;
				    
				    if (text.charAt(last-1)==',') taggedText.append(",");				    
				    if (text.charAt(last-1)==';') taggedText.append(";");				    
				    if (text.charAt(last-1)=='-') taggedText.append("-");
				    if (text.charAt(last-1)==')') taggedText.append(")");
				    if (text.charAt(last-1)=='(') taggedText.append("(");

				    start = -1;
				    type = null;
				    name.setLength(0);				    
				}
			}
		}
		taggedText.append(text.substring(last));
		return taggedText.toString();
	}	
}
